// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Common types and traits of the EPMB pallet group.
//!
//! ## [`SolutionOf`]
//!
//! This type is among the most cryptic used in the EPMB pallet. The origins of this type go back to
//! the fact that sending a solution, with hundreds or thousands of account-ids in it would be too
//! large for a chain to handle. This was particularly the case in a single page solution, as
//! developed in `election-provider-multi-phase`. To combat this, a "compact" custom type is
//! generated to encapsulate a solution. This type is generated by
//! [`frame_election_provider_support::generate_solution_type`]. See the documentation of this macro
//! for more information about the hacks used to reduce the size of the solution.
//!
//! Consequently, the [`SolutionVoterIndexOf`] and [`SolutionTargetIndexOf`] and
//! [`SolutionAccuracyOf`] are derived from this type.
//!
//! ## [`Phase`]
//!
//! This is the most important type of this pallet, demonstrating the state-machine used
//! to manage the election process and its various phases.

use crate::{unsigned::miner::MinerConfig, verifier};
use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen};
use frame_election_provider_support::ElectionProvider;
pub use frame_election_provider_support::{NposSolution, PageIndex};
use frame_support::{
	traits::DefensiveSaturating, BoundedVec, CloneNoBound, DebugNoBound, DefaultNoBound, EqNoBound,
	PartialEqNoBound,
};
use frame_system::pallet_prelude::BlockNumberFor;
use scale_info::TypeInfo;
use sp_core::Get;
pub use sp_npos_elections::{ElectionResult, ElectionScore};
use sp_runtime::{
	traits::{CheckedSub, One, Zero},
	SaturatedConversion, Saturating,
};
use sp_std::{collections::btree_set::BTreeSet, fmt::Debug, prelude::*};

/// The solution type used by this crate.
pub type SolutionOf<T> = <T as MinerConfig>::Solution;
/// The voter index. Derived from [`SolutionOf`].
pub type SolutionVoterIndexOf<T> = <SolutionOf<T> as NposSolution>::VoterIndex;
/// The target index. Derived from [`SolutionOf`].
pub type SolutionTargetIndexOf<T> = <SolutionOf<T> as NposSolution>::TargetIndex;
/// The accuracy of the election, when submitted from offchain. Derived from [`SolutionOf`].
pub type SolutionAccuracyOf<T> = <SolutionOf<T> as NposSolution>::Accuracy;
/// The fallback election type.
pub type FallbackErrorOf<T> = <<T as crate::Config>::Fallback as ElectionProvider>::Error;

/// The relative distribution of a voter's stake among the winning targets.
pub type AssignmentOf<T> =
	sp_npos_elections::Assignment<<T as MinerConfig>::AccountId, SolutionAccuracyOf<T>>;

/// A paginated raw solution type.
///
/// This is the representation of a stored, unverified solution.
///
/// After feasibility, it is converted into `Supports`.
#[derive(
	TypeInfo,
	Encode,
	Decode,
	DecodeWithMemTracking,
	DebugNoBound,
	CloneNoBound,
	EqNoBound,
	PartialEqNoBound,
	DefaultNoBound,
)]
#[codec(mel_bound(T: crate::Config))]
#[scale_info(skip_type_params(T))]
pub struct PagedRawSolution<T: MinerConfig> {
	/// The individual pages.
	pub solution_pages: Vec<SolutionOf<T>>,
	/// The final claimed score post feasibility and concatenation of all pages.
	pub score: ElectionScore,
	/// The designated round.
	pub round: u32,
}

impl<T: MinerConfig> PagedRawSolution<T> {
	/// Get the total number of voters, assuming that voters in each page are unique.
	pub fn voter_count(&self) -> usize {
		self.solution_pages
			.iter()
			.map(|page| page.voter_count())
			.fold(0usize, |acc, x| acc.saturating_add(x))
	}

	/// Get the total number of winners, assuming that there's only a single page of targets.
	pub fn winner_count_single_page_target_snapshot(&self) -> usize {
		self.solution_pages
			.iter()
			.map(|page| page.unique_targets())
			.into_iter()
			.flatten()
			.collect::<BTreeSet<_>>()
			.len()
	}

	/// Get the total number of edges.
	pub fn edge_count(&self) -> usize {
		self.solution_pages
			.iter()
			.map(|page| page.edge_count())
			.fold(0usize, |acc, x| acc.saturating_add(x))
	}
}

/// A helper trait to deal with the page index of partial solutions.
///
/// This should only be called on the `Vec<Solution>` or similar types. If the solution is *full*,
/// then it returns a normal iterator that is just mapping the index (usize) to `PageIndex`.
///
/// if the solution is partial, it shifts the indices sufficiently so that the most significant page
/// of the solution matches with the most significant page of the snapshot onchain.
///
/// See the tests below for examples.
pub trait Pagify<T> {
	/// Pagify a reference.
	fn pagify(&self, bound: PageIndex) -> Box<dyn Iterator<Item = (PageIndex, &T)> + '_>;
}

impl<T> Pagify<T> for Vec<T> {
	fn pagify(&self, desired_pages: PageIndex) -> Box<dyn Iterator<Item = (PageIndex, &T)> + '_> {
		Box::new(
			self.into_iter()
				.enumerate()
				.map(|(p, s)| (p.saturated_into::<PageIndex>(), s))
				.map(move |(p, s)| {
					let desired_pages_usize = desired_pages as usize;
					// TODO: this could be an error.
					debug_assert!(self.len() <= desired_pages_usize);
					let padding = desired_pages_usize.saturating_sub(self.len());
					let new_page = p.saturating_add(padding.saturated_into::<PageIndex>());
					(new_page, s)
				}),
		)
	}
}

/// Helper trait to pad a partial solution such that the leftover pages are filled with zero.
///
/// See the tests below for examples.
pub trait PadSolutionPages: Sized {
	/// Pad the solution to the given number of pages.
	fn pad_solution_pages(self, desired_pages: PageIndex) -> Self;
}

impl<T: Default + Clone + Debug> PadSolutionPages for Vec<T> {
	fn pad_solution_pages(self, desired_pages: PageIndex) -> Self {
		let desired_pages_usize = desired_pages as usize;
		debug_assert!(self.len() <= desired_pages_usize);
		if self.len() == desired_pages_usize {
			return self
		}

		// we basically need to prepend the list with this many items.
		let empty_slots = desired_pages_usize.saturating_sub(self.len());
		sp_std::iter::repeat(Default::default())
			.take(empty_slots)
			.chain(self.into_iter())
			.collect::<Vec<_>>()
	}
}

impl<T: Default + Clone + Debug, Bound: frame_support::traits::Get<u32>> PadSolutionPages
	for BoundedVec<T, Bound>
{
	fn pad_solution_pages(self, desired_pages: PageIndex) -> Self {
		let desired_pages_usize = (desired_pages).min(Bound::get()) as usize;
		debug_assert!(self.len() <= desired_pages_usize);
		if self.len() == desired_pages_usize {
			return self
		}

		// we basically need to prepend the list with this many items.
		let empty_slots = desired_pages_usize.saturating_sub(self.len());
		let self_as_vec = sp_std::iter::repeat(Default::default())
			.take(empty_slots)
			.chain(self.into_iter())
			.collect::<Vec<_>>();
		self_as_vec.try_into().expect("sum of both iterators has at most `desired_pages_usize` items; `desired_pages_usize` is `min`-ed by `Bound`; conversion cannot fail; qed")
	}
}

/// Alias for a voter, parameterized by the miner config.
pub type VoterOf<T> = frame_election_provider_support::Voter<
	<T as MinerConfig>::AccountId,
	<T as MinerConfig>::MaxVotesPerVoter,
>;

/// Alias for a page of voters, parameterized by this crate's config.
pub type VoterPageOf<T> = BoundedVec<VoterOf<T>, <T as MinerConfig>::VoterSnapshotPerBlock>;

/// Alias for all pages of voters, parameterized by this crate's config.
pub type AllVoterPagesOf<T> = BoundedVec<VoterPageOf<T>, <T as MinerConfig>::Pages>;

/// Maximum number of items that [`AllVoterPagesOf`] can contain, when flattened.
pub struct MaxFlattenedVoters<T: MinerConfig>(sp_std::marker::PhantomData<T>);
impl<T: MinerConfig> Get<u32> for MaxFlattenedVoters<T> {
	fn get() -> u32 {
		T::VoterSnapshotPerBlock::get().saturating_mul(T::Pages::get())
	}
}

/// Same as [`AllVoterPagesOf`], but instead of being a nested bounded vec, the entire voters are
/// flattened into one outer, unbounded `Vec` type.
///
/// This is bounded by [`MaxFlattenedVoters`].
pub type AllVoterPagesFlattenedOf<T> = BoundedVec<VoterOf<T>, MaxFlattenedVoters<T>>;

/// Current phase of the pallet.
#[derive(
	PartialEqNoBound,
	EqNoBound,
	CloneNoBound,
	Encode,
	Decode,
	DecodeWithMemTracking,
	MaxEncodedLen,
	DebugNoBound,
	TypeInfo,
)]
#[codec(mel_bound(T: crate::Config))]
#[scale_info(skip_type_params(T))]
pub enum Phase<T: crate::Config> {
	/// Nothing is happening, but it might.
	Off,
	/// Signed phase is open.
	///
	/// The inner value is the number of blocks left in this phase.
	Signed(BlockNumberFor<T>),
	/// We are validating results.
	///
	/// This always follows the signed phase, and is a window of time in which we try to validate
	/// our signed results.
	///
	/// The inner value is the number of blocks left in this phase.
	SignedValidation(BlockNumberFor<T>),
	/// Unsigned phase.
	///
	/// The inner value is the number of blocks left in this phase.
	///
	/// We do not yet check whether the unsigned phase is active or passive. The intent is for the
	/// blockchain to be able to declare: "I believe that there exists an adequate signed
	/// solution," advising validators not to bother running the unsigned offchain worker.
	///
	/// As validator nodes are free to edit their OCW code, they could simply ignore this advisory
	/// and always compute their own solution. However, by default, when the unsigned phase is
	/// passive, the offchain workers will not bother running.
	Unsigned(BlockNumberFor<T>),
	/// Snapshot is being created. No other operation is allowed. This can be one or more blocks.
	/// The inner value should be read as "`remaining` number of pages are left to be fetched".
	/// Thus, if inner value is `0` if the snapshot is complete and we are ready to move on.
	///
	/// This value should be interpreted after `on_initialize`/`on_poll` of this pallet has already
	/// been called.
	Snapshot(PageIndex),
	/// Snapshot is done, and we are waiting for `Export` to kick in.
	Done,
	/// Exporting has begun, and the given page was the last one received.
	///
	/// Once this is active, no more signed or solutions will be accepted.
	Export(PageIndex),
	/// The emergency phase. This is could be enabled by one of the fallbacks, and locks the pallet
	/// such that only governance can change the state.
	Emergency,
}

impl<T: crate::Config> Copy for Phase<T> {}

impl<T: crate::Config> Default for Phase<T> {
	fn default() -> Self {
		Phase::Off
	}
}

impl<T: crate::Config> Phase<T> {
	/// Get the phase that we should set in storage once we receive the start signal.
	pub(crate) fn start_phase() -> Self {
		// note that we add one block because we want the target snapshot to happen one block
		// before.
		Self::Snapshot(T::Pages::get())
	}

	fn are_we_done() -> Self {
		let query = T::AreWeDone::get();
		query
	}

	/// A hack to make sure we don't finish the signed verification phase just yet if the status is
	/// not yet set back to `Nothing`.
	fn verifier_done() -> bool {
		matches!(
			<T::Verifier as verifier::AsynchronousVerifier>::status(),
			verifier::Status::Nothing
		)
	}

	/// Consume self and return the next variant, as per what the current phase is.
	pub fn next(self) -> Self {
		match self {
			// for these phases, we do nothing.
			Self::Off => Self::Off,
			Self::Emergency => Self::Emergency,

			// snapshot phase
			Self::Snapshot(0) =>
				if let Some(signed_duration) = T::SignedPhase::get().checked_sub(&One::one()) {
					Self::Signed(signed_duration)
				} else if let Some(unsigned_duration) =
					T::UnsignedPhase::get().checked_sub(&One::one())
				{
					Self::Unsigned(unsigned_duration)
				} else {
					Self::are_we_done()
				},
			Self::Snapshot(non_zero_remaining) =>
				Self::Snapshot(non_zero_remaining.defensive_saturating_sub(One::one())),

			// signed phase
			Self::Signed(zero) if zero == BlockNumberFor::<T>::zero() =>
				Self::SignedValidation(T::SignedValidationPhase::get()),
			Self::Signed(non_zero_left) =>
				Self::Signed(non_zero_left.defensive_saturating_sub(One::one())),

			// signed validation
			Self::SignedValidation(zero)
				if zero == BlockNumberFor::<T>::zero() && Self::verifier_done() =>
				if let Some(unsigned_duration) = T::UnsignedPhase::get().checked_sub(&One::one()) {
					Self::Unsigned(unsigned_duration)
				} else {
					Self::are_we_done()
				},
			Self::SignedValidation(non_zero_left) =>
				Self::SignedValidation(non_zero_left.saturating_sub(One::one())),

			// unsigned phase -- at this phase we will
			Self::Unsigned(zero) if zero == BlockNumberFor::<T>::zero() => Self::are_we_done(),
			Self::Unsigned(non_zero_left) =>
				Self::Unsigned(non_zero_left.defensive_saturating_sub(One::one())),

			// Done. Wait for export to start.
			Self::Done => Self::Done,

			// Export never moves forward via this function, and is always manually set in the
			// `elect` code path.
			Self::Export(x) => Self::Export(x),
		}
	}

	/// Whether the phase is emergency or not.
	pub fn is_emergency(&self) -> bool {
		matches!(self, Phase::Emergency)
	}

	/// Whether the phase is signed or not.
	pub fn is_signed(&self) -> bool {
		matches!(self, Phase::Signed(_))
	}

	/// Whether the phase is unsigned or not.
	pub fn is_unsigned(&self) -> bool {
		matches!(self, Phase::Unsigned(_))
	}

	/// Whether the phase is off or not.
	pub fn is_off(&self) -> bool {
		matches!(self, Phase::Off)
	}

	/// Whether the phase is snapshot or not.
	pub fn is_snapshot(&self) -> bool {
		matches!(self, Phase::Snapshot(_))
	}

	/// Whether the phase is done or not.
	pub fn is_done(&self) -> bool {
		matches!(self, Phase::Done)
	}

	/// Whether the phase is export or not.
	pub fn is_export(&self) -> bool {
		matches!(self, Phase::Export(_))
	}

	/// Whether the phase is signed validation or not.
	pub fn is_signed_validation(&self) -> bool {
		matches!(self, Phase::SignedValidation(_))
	}

	/// Whether the signed phase is opened now.
	pub fn is_signed_validation_opened_now(&self) -> bool {
		self == &Phase::SignedValidation(T::SignedValidationPhase::get().saturating_sub(One::one()))
	}

	/// Whether the unsigned phase is opened now.
	pub fn is_unsigned_opened_now(&self) -> bool {
		self == &Phase::Unsigned(T::UnsignedPhase::get().saturating_sub(One::one()))
	}
}

/// Slim interface for the parent pallet to be able to inspect the state of the signed pallet.
///
/// Intentionally left different from [`crate::verifier::SolutionDataProvider`], as that is
/// specialized for communication between `verifier <> signed`.
pub trait SignedInterface {
	/// Returns `true` if there is a candidate solution to be verified.
	fn has_leader(round: u32) -> bool;
}

#[cfg(test)]
mod pagify {
	use super::{PadSolutionPages, Pagify};

	#[test]
	fn pagify_works() {
		// is a noop when you have the same length
		assert_eq!(
			vec![10, 11, 12].pagify(3).collect::<Vec<_>>(),
			vec![(0, &10), (1, &11), (2, &12)]
		);

		// pads the values otherwise
		assert_eq!(vec![10, 11].pagify(3).collect::<Vec<_>>(), vec![(1, &10), (2, &11)]);
		assert_eq!(vec![10].pagify(3).collect::<Vec<_>>(), vec![(2, &10)]);
	}

	#[test]
	fn pad_solution_pages_works() {
		// noop if the solution is complete, as with pagify.
		let solution = vec![1u32, 2, 3];
		assert_eq!(solution.pad_solution_pages(3), vec![1, 2, 3]);

		// pads the solution with default if partial..
		let solution = vec![2, 3];
		assert_eq!(solution.pad_solution_pages(3), vec![0, 2, 3]);
	}
}
